# trian.py
import json
import time

import torch
import torch.nn as nn
from ..utils.common import Common
from ..utils.data_loader import NerDataLoder
import torch.optim as optim
from ..model.BiLSTM import NERLSTM
from ..model.BiLSTM_CRF import NERLSTM_CRF
# from ..utils.data_loader import NerDataset
from tqdm import tqdm
from sklearn.metrics import precision_score, recall_score, f1_score, classification_report
from ..config.config import *



# 第二步：实现模型训练函数的搭建：mode2trian()

class LSTMModelTrainPredict():
    # 初始化参数
    def __init__(self):
        # 初始化config对象
        self.config = Config()
        self.tag2id = json.load(open(self.config.tag2id_path))
        self.common = Common()
        self.datas,self.word2id = self.common.build_data()

    def model2train(self, my_model="BiLSTM"):
        """
        实现模型训练函数的搭建
        :return:
        """
        # 获取数据
        train_dataloader, dev_dataloader = NerDataLoder().get_data()
        # 实例化模型
        models = {'BiLSTM': NERLSTM,
                  'BiLSTM_CRF': NERLSTM_CRF}
        model = models[my_model](self.config.embedding_dim, self.config.hidden_dim, self.config.dropout, self.word2id,
                                 self.tag2id)
        model = model.to(self.config.device)
        # 实例化损失函数
        criterion = nn.CrossEntropyLoss()
        # 实例化优化器
        optimizer = optim.Adam(model.parameters(), lr=self.config.lr)
        # 选择模型进行训练
        start_time = time.time()
        if my_model == 'BiLSTM':
            f1_score = -1000
            for epoch in range(self.config.epochs):
                model.train()
                for index, (inputs, labels, mask) in enumerate(tqdm(train_dataloader, desc='BiLSTM训练')):
                    x = inputs.to(self.config.device)
                    mask = mask.to(self.config.device)
                    y = labels.to(self.config.device)
                    pred = model(x, mask)
                    pred = pred.view(-1, len(self.tag2id))
                    my_loss = criterion(pred, y.view(-1))
                    optimizer.zero_grad()
                    my_loss.backward()
                    optimizer.step()
                    if index % 200 == 0:
                        print('epoch:%04d,------------loss:%f' % (epoch, my_loss.item()))
                precision, recall, f1, report = self.model2dev(dev_dataloader, model, criterion)
                if f1 > f1_score:
                    f1_score = f1
                    torch.save(model.state_dict(), f'{self.config.save_model_path}\\bilstm_best.pth')
                    print(report)
            end_time = time.time()
            print(f'训练总耗时：{end_time - start_time}')
        elif my_model == 'BiLSTM_CRF':
            f1_score = -1000
            for epoch in range(self.config.epochs):
                model.train()
                for index, (inputs, labels, mask) in enumerate(tqdm(train_dataloader, desc='bilstm+crf训练')):
                    x = inputs.to(self.config.device)
                    mask = mask.to(torch.bool).to(self.config.device)
                    tags = labels.to(self.config.device)
                    # CRF
                    loss = model.log_likelihood(x, tags, mask).mean()
                    optimizer.zero_grad()
                    loss.backward()
                    # CRF
                    torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=10)
                    optimizer.step()
                    if index % 200 == 0:
                        print('epoch:%04d,------------loss:%f' % (epoch, loss.item()))
                precision, recall, f1, report = self.model2dev(dev_dataloader, model)
                if f1 > f1_score:
                    f1_score = f1
                    torch.save(model.state_dict(), f'{self.config.save_model_path}\\bilstm_crf_best.pth')
                    print(report)
            end_time = time.time()
            print(f'训练总耗时：{end_time - start_time}')

    def model2dev(self, dev_iter, model, criterion=None):
        """
        实现模型验证函数的搭建
        :param model: 模型
        :param criterion:计算损失，默认志为None
        :return:
        """
        aver_loss = 0
        preds, golds = [], []
        model.eval()
        for index, (inputs, labels, mask) in enumerate(tqdm(dev_iter, desc="测试集验证")):
            val_x = inputs.to(self.config.device)
            mask = mask.to(self.config.device)
            val_y = labels.to(self.config.device)
            predict = []
            if model.name == "BiLSTM":
                pred = model(val_x, mask)
                predict = torch.argmax(pred, dim=-1).tolist()
                pred = pred.view(-1, len(self.config.tag2id))
                val_loss = criterion(pred, val_y.view(-1))
                aver_loss += val_loss.item()
            elif model.name == "BiLSTM_CRF":
                mask = mask.to(torch.bool)
                predict = model(val_x, mask)
                loss = model.log_likelihood(val_x, val_y, mask)
                aver_loss += loss.mean().item()
            # 统计非0的，也就是真实标签的长度
            leng = []
            for i in val_y.cpu():
                tmp = []
                for j in i:
                    if j.item() > 0:
                        tmp.append(j.item())
                leng.append(tmp)
            # 提取真实长度的预测标签

            for index, i in enumerate(predict):
                preds.extend(i[:len(leng[index])])

            # 提取真实长度的真实标签
            for index, i in enumerate(val_y.tolist()):
                golds.extend(i[:len(leng[index])])
        aver_loss /= (len(dev_iter) * 64)
        precision = precision_score(golds, preds, average='macro')
        recall = recall_score(golds, preds, average='macro')
        f1 = f1_score(golds, preds, average='macro')
        report = classification_report(golds, preds)
        return precision, recall, f1, report


        # 第二步：实现模型预测函数：model2test


    def model2test(self, sample, my_model="BiLSTM"):
            """
            实现模型预测函数
            :param self:
            :param sample:
            :return:
            """
            # 实例化模型
            models = {'BiLSTM': NERLSTM,
                      'BiLSTM_CRF': NERLSTM_CRF}
            model = models[my_model](self.config.embedding_dim, self.config.hidden_dim, self.config.dropout,
                                         self.word2id, self.tag2id)
            if my_model == 'BiLSTM':
                model.load_state_dict(torch.load(f'{self.config.save_model_path}\\bilstm_best.pth'))
            else:
                model.load_state_dict(torch.load(f'{self.config.save_model_path}\\bilstm_crf_best.pth'))
            id2tag = {value: key for key, value in self.tag2id.items()}

            x = []
            for char in sample:
                if char not in self.word2id:
                    char = "UNK"
                x.append(self.word2id[char])

            x_train = torch.tensor([x])
            mask = (x_train != 0).long()
            model.eval()
            with torch.no_grad():
                if model.name == "BiLSTM":
                    outputs = model(x_train, mask)
                    preds_ids = torch.argmax(outputs, dim=-1)[0]
                    tags = [id2tag[i.item()] for i in preds_ids]
                else:
                    preds_ids = model(x_train, mask)
                    tags = [id2tag[i] for i in preds_ids[0]]
                chars = [i for i in sample]
                assert len(chars) == len(tags)
                result = self.extract_entities(chars, tags)
                return result
    def extract_entities(self, tokens, labels):
            
            """
            实现实体解析函数
            :param tokens: 文本的所有字符
            :param labels: 标签
            :return:
            """
            entities = []
            entity = []
            entity_type = None
            for token, label in zip(tokens, labels):
                if label.startswith("B-"):  # 实体的开始
                    if entity:  # 如果已经有实体，先保存
                        entities.append((entity_type, ''.join(entity)))
                        entity = []
                    entity_type = label.split('-')[1]
                    entity.append(token)
                elif label.startswith("I-") and entity:  # 实体的中间或结尾
                    entity.append(token)
                else:
                    if entity:  # 保存上一个实体
                        entities.append((entity_type, ''.join(entity)))
                        entity = []
                        entity_type = None

            # 如果最后一个实体没有保存，手动保存
            if entity:
                entities.append((entity_type, ''.join(entity)))

            return {entity: entity_type for entity_type, entity in entities}
    
    def model_application(self, logo, text=""):
        """
        logo : 标识符,判断是train表示对模型进行训练,test表示对模型进行测试
        """
        # 实例化对象
        if logo == "train":
            self.model2train()
        elif logo == "predict":
            # 模型预测
            result = self.model2test(text,"BiLSTM")
            return result
            